package main

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
)

// 全局SSH客户端连接和互斥锁
var client *ssh.Client
var clientMutex sync.Mutex

// 定义一个特殊的字符串，用于标记命令输出的结束
const outputDelimiter = "[END_OF_OUTPUT]"

func main() {
	fmt.Println("SSH Console v1.0 - Ready")
	fmt.Println("Listening for commands on stdin...")

	// 使用bufio.Scanner读取标准输入，按行处理
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		line := scanner.Text()
		parts := strings.Fields(line)
		if len(parts) == 0 {
			continue
		}

		command := parts[0]
		args := parts[1:]

		switch command {
		case "CONNECT":
			handleConnect(args)
		case "COMMAND":
			handleCommand(args)
		case "DISCONNECT":
			handleDisconnect()
		case "EXIT":
			handleDisconnect()
			return // 退出主循环，结束程序
		default:
			fmt.Printf("ERROR: Unknown command: %s\n%s\n", command, outputDelimiter)
		}
	}

	if err := scanner.Err(); err != nil {
		fmt.Printf("ERROR: Reading stdin failed: %v\n%s\n", err, outputDelimiter)
	}
}

// handleConnect 处理 CONNECT 命令
// 格式: CONNECT <host> <port> <user> <password>
func handleConnect(args []string) {
	if len(args) != 4 {
		fmt.Printf("ERROR: Invalid arguments for CONNECT. Usage: CONNECT <host> <port> <user> <password>\n%s\n", outputDelimiter)
		return
	}

	clientMutex.Lock()
	defer clientMutex.Unlock()

	// 如果已连接，先断开
	if client != nil {
		client.Close()
		client = nil
	}

	host := args[0]
	port := args[1]
	user := args[2]
	password := args[3]
	addr := fmt.Sprintf("%s:%s", host, port)

	cfg := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil // 接受所有主机密钥，简化处理
		},
		Timeout: 5 * time.Second,
	}

	var err error
	client, err = ssh.Dial("tcp", addr, cfg)
	if err != nil {
		fmt.Printf("ERROR: Failed to dial: %v\n%s\n", err, outputDelimiter)
		return
	}

	fmt.Printf("SUCCESS: Connected to %s\n%s\n", addr, outputDelimiter)
}

// handleCommand 处理 COMMAND 命令
// 格式: COMMAND <command_string>
func handleCommand(args []string) {
	clientMutex.Lock()
	defer clientMutex.Unlock()

	if client == nil {
		fmt.Printf("ERROR: Not connected. Please use CONNECT first.\n%s\n", outputDelimiter)
		return
	}

	// 重新组合命令字符串
	command := strings.Join(args, " ")

	session, err := client.NewSession()
	if err != nil {
		fmt.Printf("ERROR: Failed to create session: %v\n%s\n", err, outputDelimiter)
		return
	}
	defer session.Close()

	// 运行命令并将合并的输出打印到stdout
	output, err := session.CombinedOutput(command)
	if err != nil {
		fmt.Printf("ERROR: Failed to run command: %v\nOutput: %s\n%s\n", err, string(output), outputDelimiter)
		return
	}

	// 将输出直接写入标准输出
	fmt.Print(string(output))
	// 打印分隔符，告知易语言模块命令执行完毕
	fmt.Printf("%s\n", outputDelimiter)
}

// handleDisconnect 处理 DISCONNECT 命令
func handleDisconnect() {
	clientMutex.Lock()
	defer clientMutex.Unlock()

	if client != nil {
		client.Close()
		client = nil
		fmt.Printf("SUCCESS: Disconnected\n%s\n", outputDelimiter)
	} else {
		fmt.Printf("SUCCESS: Already disconnected\n%s\n", outputDelimiter)
	}
}
